from rpyc.utils.security.base import BaseDB, Channel
from rpyc.utils.security.base import ProtocolError, ProtocolAuthError
from rpyc.utils.security.crypto_integration import RSAKey, BadRSAKey, get_random_bytes


#
# public-key database
#
class PubkeyDB(BaseDB):
    """public-key database (pkdb), in SSH terminology, is basically a unified
    container of an id_rsa file alongside with an authorized_keys file.
    the first line is the private key of the client or server, followed by
    "trusted public keys", i.e., for a client that would be known_hosts, and
    for the server, authorized_keys. 
    
    the pkdb allows the two parties to mutually authenticate themselves.
    """
    RSA_BITS = 2048
    __slots__ = ["private_key", "_pub_keys_by_label", "_pub_keys_by_key"]
    
    def __init__(self, private_key, known_pub_keys, filename = None):
        BaseDB.__init__(self, filename)
        self.private_key = private_key
        self._pub_keys_by_label = {}
        self._pub_keys_by_key = {}
        for label, key in known_pub_keys.iteritems():
            self.add_public_key(label, key)
    
    def __len__(self):
        return len(self._pub_keys_by_label)
    def __iter__(self):
        return self._pub_keys_by_label.iteritems()
    def __getitem__(self, label):
        return self._pub_keys_by_label[label]
    def get(self, label, default = None):
        try:
            return self[label]
        except KeyError:
            return default
    def __contains__(self, label):
        return label in self._pub_keys_by_label

    @classmethod
    def create(cls, filename = None):
        return cls(RSAKey.generate(cls.RSA_BITS), {}, filename)
    
    @classmethod
    def _load(cls, data):
        inst = cls(None, {})
        for line in data.splitlines():
            line = line.strip()
            if not line:
                continue # empty line
            if line.startswith("#"):
                continue # comment line
            tokens = line.split(" ", 1)
            if len(tokens) == 1:
                rawkey, label = tokens[0], ""
            else:
                rawkey, label = tokens
            key = RSAKey.load(rawkey)
            if not inst.private_key:
                inst.private_key = key
                continue # the first key is the private key
            inst.add_public_key(label.strip(), key)
        if not inst.private_key:
            raise ValueError("no keys found; at least the private_key is required")
        return inst
    
    def _dump(self):
        lines = ["# PRIVATE KEY",
            self.private_key.dump_private(),
            "# KNOWN PUBLIC KEYS",
        ]
        for label, pk in self:
            lines.append("%s %s" % (pk.dump_public(), label))
        return "\n".join(lines) + "\n"

    def add_public_key(self, label, pubkey):
        if label in self._pub_keys_by_label:
            raise ValueError("label %r already exists" % (label,))
        self._pub_keys_by_label[label] = pubkey
        if pubkey not in self._pub_keys_by_key:
            self._pub_keys_by_key[pubkey] = set()
        self._pub_keys_by_key[pubkey].add(label)
    
    def remove_public_key(self, label):
        pubkey = self._pub_keys_by_label.pop(label)
        self._pub_keys_by_key[pubkey].remove(label)
        if not self._pub_keys_by_key[pubkey]:
            del self._pub_keys_by_key[pubkey]
    
    def remove_public_key_by_key(self, pubkey):
        for label in list(self._pub_keys_by_key[pubkey]):
            self.remove_public_key(label)

    def contains_public_key(self, pubkey):
        return pubkey in self._pub_keys_by_key


#
# public-key authentication protocol
#
class PubkeyProtocol(object):
    CHALLENGE_BYTES = 32
    
    @classmethod
    def _initiate_handshake(cls, chan, pubkey):
        blob = get_random_bytes(cls.CHALLENGE_BYTES)
        chan.send(blob)
        sig = chan.recv()
        if pubkey.verify(blob, sig):
            chan.send("OK")
        else:
            chan.send("ERROR")
            raise ProtocolAuthError("wrong signature")
    
    @classmethod
    def _respond_handshake(cls, chan, prikey):
        blob = chan.recv()
        sig = prikey.sign(blob)
        chan.send(sig)
        succ = chan.recv()
        if succ != "OK":
            raise ProtocolAuthError("signature rejected")
    
    @classmethod
    def handshake_server(cls, sock, pubkeydb):
        chan = Channel(sock)
        chan.send("PUBKEY")
        raw_client_pubkey = chan.recv()
        client_pubkey = RSAKey.load(raw_client_pubkey)
        if not pubkeydb.contains_public_key(client_pubkey):
            client_pubkey = BadRSAKey
        cls._initiate_handshake(chan, client_pubkey)
        cls._respond_handshake(chan, pubkeydb.private_key)
        return sock, client_pubkey
    
    @classmethod
    def handshake_client(cls, sock, private_key, server_pubkey):
        chan = Channel(sock)
        chan.expect("PUBKEY")
        chan.send(private_key.dump_public())
        cls._respond_handshake(chan, private_key)
        cls._initiate_handshake(chan, server_pubkey)
        return sock

